import numpy as np
import pandas as pd
from qtorch.strategy import Strategy

class MeanReversionStrategy(Strategy):
    """均值回归策略（20日Z-Score）"""
    def __init__(self, window=20, entry_z=1.0, exit_z=0.5):
        self.window = window
        self.entry_z = entry_z
        self.exit_z = exit_z
        
    def generate_signals(self, prices):
        rolling_mean = prices.rolling(window=self.window).mean()
        rolling_std = prices.rolling(window=self.window).std()
        z_scores = (prices - rolling_mean) / rolling_std
        
        signals = np.where(z_scores > self.entry_z, -1,
                         np.where(z_scores < -self.entry_z, 1,
                                np.where(abs(z_scores) < self.exit_z, 0, np.nan)))
        signals = pd.Series(signals).ffill().values
        return signals.astype(int)